//
//  BinaryNode.swift
//  BinaryTree
//
//  Created by 郭红旗 on 2021/12/24.
//

import UIKit

class BinaryBRNodeManager<B:Equatable>: BinarySearchNodeManager<B> {
    
    override func printFun() {
        print(root!.asRBString)
    }
    
    override func afterAdd(node: BinaryNode<B>) {
        // 添加的根节点
        guard let parent = node.parent else {
            balckNode(node: node)
            return
        }
        // 父节点是黑色直接返回
        if isBalck(node: parent) {return}
        guard let grand = parent.parent else {
            return
        }
        // 父节点为红色
        // 父节点的兄弟节点存在，并且是红色
        if let uncle = parent.sibling() {
            if isRed(node: uncle) {
                balckNode(node: uncle)
                balckNode(node: parent)
                redNode(node: grand)
                afterAdd(node: grand)
                return
            }
        }
        // 父节点不存在兄弟节点
        if parent.isLeft() { // L
            if node.isLeft() {// LL
                redNode(node: grand)
                balckNode(node: parent)
            }else{ // LR
                balckNode(node: node)
                redNode(node: grand)
                rotateLeft(tree: parent)
            }
            rotateRight(tree: grand)

        }else{ // R
            if node.isRight() {// RR
                redNode(node: grand)
                balckNode(node: parent)
            }else{ // RL
                balckNode(node: node)
                redNode(node: grand)
                rotateRight(tree: parent)
            }
            rotateLeft(tree: grand)
        }
    }
    
    /// 右旋
    override func rotateRight(tree:BinaryNode<B>) {
      super.rotateRight(tree: tree)
    }
    /// 左旋
    override func rotateLeft(tree:BinaryNode<B>) {
      super.rotateLeft(tree: tree)
    }
    
    func balckNode(node:BinaryNode<B>) {
        node.color = "balck"
    }
    func redNode(node:BinaryNode<B>) {
        node.color = "red"
    }
    func isBalck(node:BinaryNode<B>) -> Bool {
        if node.color == "balck" {
            return true
        }
        return false
    }
    func isRed(node:BinaryNode<B>) -> Bool {
        if node.color == "red" {
            return true
        }
        return false
    }
}

/// 平衡树
class BinaryAVLNodeManager<A:Equatable>: BinarySearchNodeManager<A>{
    override func afterAdd(node: BinaryNode<A>) {
        super.afterAdd(node: node)
        var tree:BinaryNode<A>? = node
        while tree != nil {
            if isBalance(node: tree!) {
                updateHight(node: tree!)
            }else{
                balanceFunc(node: tree!)
                return
            }
            tree = tree!.parent
        }
    }
    
    override func afterRemove(node: BinaryNode<A>) {
        super.afterRemove(node: node)
        var tree = node.parent
        while tree != nil {
            if isBalance(node: tree!) {
                updateHight(node: tree!)
            }else{
                balanceFunc(node: tree!)
            }
            tree = tree?.parent
        }
    }
    
    func isBalance(node:BinaryNode<A>) -> Bool {
        return node.isBalance()
    }
    func updateHight(node:BinaryNode<A>){
        node.updateHight()
    }
    
    /// 恢复平衡
    func balanceFunc(node:BinaryNode<A>){
        let grand = node
        guard let parent = grand.trailChild(),let child = parent.trailChild() else {
            return
        }
        
        if parent.isLeft() { // L
            if child.isLeft() { // LL
                rotateRight(tree: grand)
            }else{ // LR
                rotateLeft(tree: parent)
                rotateRight(tree: grand)
            }
        }else{// R
            if child.isLeft() { // RL
                rotateRight(tree: parent)
                rotateLeft(tree: grand)
            }else{ // RR
                rotateLeft(tree: grand)
            }
        }
    }
    /// 右旋
    override func rotateRight(tree:BinaryNode<A>) {
        let grand = tree
        let parent = grand.left
        let child = parent?.right
        
        grand.left = child
        parent?.right = grand
        parent?.parent = grand.parent

        if grand.isLeft() {
            grand.parent?.left = parent
        }else if grand.isRight() {
            grand.parent?.right = parent
        }else{
            root = parent
        }
        
        child?.parent = grand
        grand.parent = parent
        updateHight(node: grand)
        updateHight(node: parent!)

    }
    /// 左旋
    override func rotateLeft(tree:BinaryNode<A>) {
        let grand = tree
        let parent = grand.right
        let child = parent?.left
        
        grand.right = child
        parent?.left = grand
        parent?.parent = grand.parent

        if grand.isLeft() {
            grand.parent?.left = parent
        }else if grand.isRight() {
            grand.parent?.right = parent
        }else{
            root = parent
        }
        child?.parent = grand
        grand.parent = parent
        updateHight(node: grand)
        updateHight(node: parent!)
    }
    
}

/// 二叉搜索树
class BinarySearchNodeManager<U:Equatable>{
    typealias comparatorClosure = (_ e1:U,_ e2:U) -> Int
    /// 比较器
    private var comparator:comparatorClosure
   /// 根节点
    var root:BinaryNode<U>?
    /// 初始化
    init(compare:@escaping comparatorClosure) {
        comparator = compare
    }
    /// 打印函数
    func printFun(){
        print(root!.asString)
    }
    /// 添加节点
    func addNode(withNonRecursive value:U){
        /// 根节点为空
        guard let rootNode = root else {
            root = BinaryNode.init(value: value)
            afterAdd(node: root!)
            return
        }
        var node:BinaryNode? = rootNode
        var last = node
        var comparatorResult = 0
        /// 遍历获取叶子结点
        while node != nil {
            let result = comparator(value,node!.element!)
            last = node
            comparatorResult = result
            if result > 0 {
                node = node!.right
            }else if result < 0{
                node = node!.left
            }else{
                
            }
        }
        /// 添加节点
        let tree = BinaryNode.init(value: value,parent: last)
        if comparatorResult > 0 {
            last?.right = tree
        }else if comparatorResult < 0{
            last?.left = tree
        }else{
            
        }
        afterAdd(node: tree)
    }
    
    func removeNode(withNonRecursive value:U) {
        guard var node = node(withValue: value) else {
            return
        }
        if node.hasTwoChild() {
            let replace = successorNode(node: node)
            node.element = replace.element
            node = replace
        }
        
        if node.isLeafNode() {
            if node.isLeft() {
                node.parent?.left = nil
            }else{
                node.parent?.right = nil
            }
            afterRemove(node: node)
        }else{
            let replace = node.left == nil ? node.right : node.left
            if node.isLeft() {
                node.parent?.left = replace
            }else{
                node.parent?.right = replace
            }
            replace?.parent = node.parent
            afterRemove(node: replace!)
        }
        
        
    }
    /// 添加节点 递归
    func addNode(withRecursive node:BinaryNode<U>?,value:U) -> BinaryNode<U>{
        guard let tree = node else {
            let newTree = BinaryNode.init(value: value)
            return newTree
        }
        let result = comparator(value,tree.element!)
        if result > 0 {
            tree.right = addNode(withRecursive: tree.right, value: value)
        }else{
            tree.left = addNode(withRecursive: tree.left, value: value)
        }
        return tree
    }
    func afterAdd(node:BinaryNode<U>){}
    func afterRemove(node:BinaryNode<U>){}
    func node(withValue value:U) -> BinaryNode<U>? {
        var node = root
        while node != nil {
            let result = comparator(value,node!.element!)
            if result > 0 {
                node = node?.right
            }else if result < 0{
                node = node?.left
            }else{
                return node
            }
        }
        return nil
    }
    // 前驱节点
    func forerunnerNode(node:BinaryNode<U>) -> BinaryNode<U>{
        var tree = node.left
        var result = tree
        while tree != nil {
            result = tree
            tree = tree?.right
        }
        return result!
    }
    // 后继节点
    func successorNode(node:BinaryNode<U>) -> BinaryNode<U>{
        var tree = node.right
        var result = tree
        while tree != nil {
            result = tree
            tree = tree?.left
        }
        return result!
    }
    /// 右旋
    func rotateRight(tree:BinaryNode<U>) {
        let grand = tree
        let parent = grand.left
        let child = parent?.right
        
        grand.left = child
        parent?.right = grand
        parent?.parent = grand.parent

        if grand.isLeft() {
            grand.parent?.left = parent
        }else if grand.isRight() {
            grand.parent?.right = parent
        }else{
            root = parent
        }
        
        child?.parent = grand
        grand.parent = parent


    }
    /// 左旋
    func rotateLeft(tree:BinaryNode<U>) {
        let grand = tree
        let parent = grand.right
        let child = parent?.left
        
        grand.right = child
        parent?.left = grand
        parent?.parent = grand.parent

        if grand.isLeft() {
            grand.parent?.left = parent
        }else if grand.isRight() {
            grand.parent?.right = parent
        }else{
            root = parent
        }
        child?.parent = grand
        grand.parent = parent

    }
    
    // MARK: -翻转
    func flipTree(tree:BinaryNode<U>?) -> BinaryNode<U>?{
        guard let node = tree else {
            return nil
        }
        let left = flipTree(tree: node.left)
        let right = flipTree(tree: node.right)
        node.left =  right
        node.right = left
        return node
    }
    // MARK: -前序  根-左-右
    /// 递归
    func preorderTraversal(tree:BinaryNode<U>?,callBack:(U?)->Void){
        guard let node = tree else {
            return
        }
        callBack(node.element)
        preorderTraversal(tree: node.left, callBack: callBack)
        preorderTraversal(tree: node.right, callBack: callBack)
    }

    /// 非递归
    func preorderTraversalNonrecursion(tree:BinaryNode<U>) {
        var result:[BinaryNode<U>] = []
        var node:BinaryNode<U>? = tree
        while node != nil || result.count != 0 {
            if node != nil {
                result.append(node!)
                print(node!.element!)
                node = node?.left
            }else{
                guard let last = result.last else {
                    continue
                }
                node = last.right
                result.removeLast()
            }
        }
    }
    
    // MARK: -中序  左-根-右
    func inorderTraversal(tree:BinaryNode<U>?,callBack:(U?)->Void){
        guard let node = tree else {
            return
        }
        preorderTraversal(tree: node.left, callBack: callBack)
        callBack(node.element)
        preorderTraversal(tree: node.right, callBack: callBack)
    }
    /// 非递归
    func inorderTraversalNonrecursion(tree:BinaryNode<U>) {
        var result:[BinaryNode<U>] = []
        var node:BinaryNode<U>? = tree
        while node != nil || result.count != 0 {
            if node != nil {
                result.append(node!)
                node = node?.left
            }else{
                guard let last = result.last else {
                    continue
                }
                print(last.element!)
                node = last.right
                result.removeLast()
            }
        }
    }
    
    // MARK: -后序   左-右-根
    func postorderTraversal(tree:BinaryNode<U>?,callBack:(U?)->Void){
        guard let node = tree else {
            return
        }
        postorderTraversal(tree: node.left, callBack: callBack)
        postorderTraversal(tree: node.right, callBack: callBack)
        callBack(node.element)
    }
    /// 非递归
    func postorderTraversalNonrecursion(tree:BinaryNode<U>) {
        var result:[BinaryNode<U>] = []
        var p:BinaryNode<U>? = tree
        var node:BinaryNode<U>? = tree

        while p != nil {
            while node?.left != nil {
                result.append(node!)
                node = node!.left
            }
            while node != nil && (node?.right == nil || node?.right?.element == p?.element) {
                print(node?.element)
                p = node
                if result.count == 0 {
                    return
                }
                node = result.last
                result.removeLast()
            }
            guard let newNode = node else {
                return
            }
            result.append(newNode)
            node = newNode.right
        }
    }
//    func postorderTraversal(node:inout BinaryNode<U>){
//        var stack:[BinaryTreeModel] = []
//        var p = node
//        while p != nil {
//            while node.leftNode != nil {
//                stack.append(node)
//                node = node?.leftNode
//            }
//            while node != nil && (node?.rightNode == nil || node?.rightNode == p) {
//                print(node?.value ?? "")
//                p = node
//                if stack.count == 0 {
//                    return
//                }
//                node = stack.last!
//                stack.removeLast()
//            }
//            stack.append(node)
//            node = node?.rightNode
//        }
//    }
}

/// 二叉树
class BinaryNode<T:Equatable>{
    
    var left:BinaryNode?
    var right:BinaryNode?
    var parent:BinaryNode?
    var element:T?
    var height:Int = 0
    var color:String = "red"
    /// 打印函数
    var asString:String {return treeString(self){("\($0.element!)",$0.left,$0.right)}}
    var asRBString:String {return treeString(self){("\($0.color)_\($0.element!)",$0.left,$0.right)}}

    init(value:T,parent:BinaryNode? = nil) {
        self.element = value
        self.parent = parent
    }
    
    func updateHight(){
        let leftHight = left?.height ?? 0
        let rightHight = right?.height ?? 0
        height = (leftHight > rightHight ? leftHight : rightHight) + 1
    }

    func isBalance() -> Bool {
        let leftHight = left?.height ?? 0
        let rightHight = right?.height ?? 0
        let result = abs(leftHight - rightHight)
        if result > 1  {
            return false
        }
        return true
    }
    
    func trailChild() -> BinaryNode<T>? {
        let leftHight = left?.height ?? 0
        let rightHight = right?.height ?? 0
        if leftHight > rightHight {
            return left!
        }
        if rightHight > leftHight {
            return right!
        }
        if isLeft() {
            return left
        }else{
            return right
        }
    }
    
    func isLeft() -> Bool {
        if element == parent?.left?.element {
            return true
        }
        return false
    }
    func isRight() -> Bool {
        if element == parent?.right?.element {
            return true
        }
        return false
    }
    
    func hasTwoChild() -> Bool {
        if left != nil && right != nil {
            return true
        }
        return false
    }
   
    func isLeafNode() -> Bool {
        if left == nil && right == nil {
            return true
        }
        return false
    }
    /// 兄弟节点
    func sibling() -> BinaryNode<T>?{
        if isLeft() {
            return parent?.right
        }
        if isRight() {
            return parent?.left
        }
        return nil
    }
}


